from langchain_ollama import OllamaLLM # 用于基础LLM
from langchain_ollama  import ChatOllama  # 用于聊天模型

import asyncio
import yaml

import sys
sys.path.append(r"G:\system\desktop\WeChatBot")

from config.PathConfig import OLLAMA_CONFIG
from langchain_core.messages import HumanMessage, SystemMessage

class LangchainClient:
    def __init__(self,tools = True):
        # 从 yaml 文件加载配置
        with open(OLLAMA_CONFIG, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
        if tools :
            self.model = config["ollama"]["QWEN_MODEL"]
        else:
            self.model = config["ollama"]["DS_MODEL"]
        self.host = config["ollama"]["host"]
        self.client = ChatOllama(   
            base_url = self.host,
            model = self.model,  # 替换为你的模型名称
            temperature=0.7,
            max_tokens=1024,
            num_gpu=1,  # 使用GPU加速
            num_thread=8,  # CPU线程数
            repeat_penalty=1.2,  # 抑制重复内容
            
        )
    
    async def async_chat(self, messages, tools=None):
            """
            异步聊天方法
            :param messages: 消息列表，支持 LangChain 的消息格式
            :param tools: 可选的工具列表
            :return: 模型响应
            """
            # 确保 messages 是 LangChain 兼容的格式
            # 如果有工具，转换为工具调用的格式
            if tools:
                # 使用 bind_tools 方法绑定工具（LangChain 最新支持方式）
                response = await self.client.bind_tools([tools]).ainvoke(messages)
            else:
                # 无工具时直接调用
                response = await self.client.ainvoke(messages)

            return response
    
    
# async def main():
#     # 1. 初始化客户端（无工具模式）
#     client = LangchainClient(tools=False)
    
#     # 2. 构造对话消息（支持两种格式）
#     messages = [
#         # 方式一：直接使用字典格式
#         {"role": "system", "content": "你是一个专业助理"},
#         {"role": "user", "content": "量子计算的主要优势是什么？"},
        
#         # 方式二：使用 LangChain 原生消息类型
#         # SystemMessage(content="你是一个专业助理"),
#         # HumanMessage(content="量子计算的主要优势是什么？")
#     ]
#     # 3. 发起异步请求
#     response = await client.async_chat(messages)
    
#     # 4. 解析响应结果
#     print("回复内容:", response.content)
#     print("元数据:", response.response_metadata)

# if __name__ == "__main__":
#     asyncio.run(main())


# # <!--IMPORTS:[{"imported": "ChatPromptTemplate", "source": "langchain_core.prompts", "docs": "https://python.langchain.com/api_reference/core/prompts/langchain_core.prompts.chat.ChatPromptTemplate.html", "title": "OllamaLLM"}, {"imported": "OllamaLLM", "source": "langchain_ollama.llms", "docs": "https://python.langchain.com/api_reference/ollama/llms/langchain_ollama.llms.OllamaLLM.html", "title": "OllamaLLM"}]-->
# from langchain_core.prompts import ChatPromptTemplate
# # from langchain_ollama.llms import OllamaLLM

# template = """Question: {question}

# Answer: Let's think step by step."""

# prompt = ChatPromptTemplate.from_template(template)

# model = OllamaLLM(model="qwen2.5:7b",base_url='http://localhost:11434')

# chain = prompt | model

# print(chain.invoke({"question": "What is LangChain?"}))